﻿namespace LambdaCalc.Spike

open System

[<AutoOpen>]
module TermOperations =

    let rec identifiers = function
        | Id x -> Set.singleton x
        | Lambda (x, body) -> Set.add x (identifiers body)
        | App (f, a) -> Set.union (identifiers f) (identifiers a)

    let rec rename (newY, forX) = function
        | Id x when x = forX -> Id newY
        | Id _ as z -> z
        | App(m,n) -> App (rename (newY, forX) m, rename (newY, forX) n)
        | Lambda (x,m) when x = forX -> Lambda (newY, rename (newY, forX) m)
        | Lambda (z,m) -> Lambda(z, rename (newY, forX) m)

    let normalAlphaRename term =
        let id = ref 0
        let next() = 
            incr id
            sprintf "%%%s" ((!id).ToString())
        let rec inner = function
            | Id x -> 
                if x.StartsWith("%")
                then Id x
                else Id (next())
            | Lambda (x, body) ->
                let newArg = next()
                let body' = rename (newArg, x) body
                Lambda (newArg, inner body')
            | App (f, arg) ->
                App (inner f, inner arg)
        inner term

    let equals t1 t2 = (normalAlphaRename t1) = (normalAlphaRename t2)

    let rec freeVariables = function
        | Id x -> Set.singleton x
        | Lambda (x, body) ->
            Set.remove x (freeVariables body)
        | App (f, arg) ->
            Set.union (freeVariables f) (freeVariables arg)

    let rec boundVariables = function 
        | Id _ -> Set.empty
        | Lambda (x, body) ->
            Set.add x (boundVariables body)
        | App (f, arg) ->
            Set.union (boundVariables f) (boundVariables arg)

    let rec findFreshVar freeVars (x : Identifier) =
        if freeVars |> Set.contains x
        then findFreshVar freeVars (x + "\'")
        else x

    let rec substituteFree (varName, withExpression) = function
        | Id x when x = varName -> 
            withExpression
        | Lambda (x, body) when x <> varName ->
            let fv = freeVariables withExpression
            if fv |> Set.contains x
            then 
                let freshV = findFreshVar fv x
                let body' = rename (freshV, x) body |> substituteFree (varName, withExpression)
                Lambda(freshV, body')
            else
                Lambda(x, substituteFree (varName, withExpression) body)
        | App (f, arg) -> 
            App (
                substituteFree (varName, withExpression) f,
                substituteFree (varName, withExpression) arg)
        | expr -> expr

    let rec normalOrderBetaReduction = function
        | App (Lambda(n, App (e, Id n')), a) when n = n' -> App(e, a) 
        | App(f, a) ->
            match normalOrderBetaReduction f with
            | Lambda (n, b) ->
                b |> substituteFree (n, a) 
                  |> normalOrderBetaReduction
            | f' ->
                let a' = normalOrderBetaReduction a
                App (f, a')
        | Lambda (x, body) ->
            Lambda (x, body |> normalOrderBetaReduction)
        | exp -> exp

    let solve = normalOrderBetaReduction

    let (<.) toTerm exp =
        App(toTerm, exp) |> solve

    let (.>) exp toTerm = toTerm <. exp 